from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model
from transformers import TrainingArguments
from datasets import load_dataset
from transformers import Trainer
from peft import PeftModel
from transformers import pipeline
import torch
import pandas as pd
from datasets import Dataset,concatenate_datasets



"""
训练数据为json或csv
{"text": "用户：你好\n助手：您好，我是AI助手"}
转换格式为：
messages = [
    # 对话1：单轮问答
    [
        {"role": "user", "content": "什么是人工智能？"},
        {"role": "assistant", "content": "人工智能是研究如何使计算机模拟人类智能的科学与技术。"}
    ],
    # 对话2：多轮对话
    [
        {"role": "user", "content": "推荐一本机器学习的入门书籍。"},
        {"role": "assistant", "content": "《机器学习实战》是一本不错的入门选择。"},
        {"role": "user", "content": "这本书适合完全没有编程基础的人吗？"},
        {"role": "assistant", "content": "不太适合，建议先学习Python基础再阅读。"}
    ]
]


"""

def peft_model():
    model_name = "Qwen/Qwen3-1.7B"
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)   # 加载分词器
    tokenizer.pad_token = tokenizer.eos_token  # 设置padding符号

    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        device_map="auto",
        trust_remote_code=True,
        torch_dtype=torch.bfloat16

    )
 
    lora_config = LoraConfig(
    r=8,                  # 低秩矩阵的秩（越小参数量越少）
    lora_alpha=32,        # 缩放因子（影响学习率）
    target_modules=["q_proj", "v_proj"],  # 目标模块（通常为注意力层的Q/V矩阵）
    lora_dropout=0.1,     # Dropout率防止过拟合
    bias="none",          # 不训练偏置项
    task_type="CAUSAL_LM" # 因果语言模型任务
    )                               

    peft_model = get_peft_model(model, lora_config)   # 加载PEFT模型

    peft_model.print_trainable_parameters()  # 查看可训练参数占比（通常<1%）

    # dataset = load_dataset("csv", data_files="your_data.csv")  # 替换为实际路径
    # 1. 加载CSV数据
    def load_csv_data(file_path):
        df = pd.read_csv(file_path, sep='_!_',header=None)
        # 假设CSV包含'input'和'output'两列（Alpaca格式）
        df.columns = ['input', 'output']
        
        return Dataset.from_pandas(df)

    #dataset要适配连续对话场景，需要列表嵌套
    dataset = []
    dataset_sub = load_csv_data(r"E:\python-knowledge\6项目实战\大模型微调\toutiao_cat_data\qwen.txt")  # 替换为你的CSV路径
    dataset.append(dataset_sub)

    # 2. 数据预处理
    def preprocess(examples):       # 传进来的examples是所有训练数据，不是一个一个的传

        """
        一个examples就是一个多轮对话，一个多轮对话有多个user和assistant

        """
        MAX_LENGTH = 2048  # 根据显存调整

        #batched=True时，传进来的examples是列表，列表中嵌套的是列表，列表中嵌套的是字典，格式处理  
        # print(examples)       # 查看examples的格式

        # 多轮对话
        # messages=[]    
        # input_list = examples["input"]
        # output_list = examples["output"]
        # for index,example in enumerate(input_list):
        #     messages.append({"role": "user", "content": example})
        #     messages.append({"role": "assistant", "content": output_list[index]})
        # print(messages)

        # 单轮对话
        messages=[]    
        input_list = examples["input"]
        output_list = examples["output"]
        for index,example in enumerate(input_list):
            messages.append([{"role": "user", "content": example},{"role": "assistant", "content": output_list[index]}])

        """
            tokenized = tokenizer.apply_chat_template(
            messages,
            tokenize=True,      #tokenize=True,返回 Token ID 序列（即分词后的数值列表）,False返回对话字符串
            truncation=True,
            max_length=MAX_LENGTH,
            padding="max_length",
            add_generation_prompt=False,
            return_tensors="pt"      
        )   
            
        transformers 库在 4.30 之后对 apply_chat_template 进行了优化，但部分旧版本或第三方模型的 tokenizer 可能未同步更新，导致批量处理时不返回字典。
        关键修复逻辑
        通过 “先单条处理，再手动批量拼接” 的方式，强制生成包含 input_ids 和 attention_mask 的字典：
        """


        # 1. 先对每条数据单独处理，得到列表形式的input_ids

        #单轮对话
        single_inputs_list = [
            tokenizer.apply_chat_template(
                dialog,
                tokenize=True,
                add_generation_prompt=False,
                return_tensors=None  # 不返回张量，返回列表
            )
            for dialog in messages                    #分成单个因果问题，每个问题都是一个列表

        ]

        #多轮对话
        # single_inputs_list=tokenizer.apply_chat_template(
        #         messages,                                #连续的问答，合成一个列表
        #         tokenize=True,
        #         add_generation_prompt=False,
        #         return_tensors=None  # 不返回张量，返回列表
        #     )    

        # 2. 用tokenizer.pad批量处理，强制返回包含input_ids和attention_mask的字典
        batch_inputs = tokenizer.pad(
            {"input_ids": single_inputs_list},         #一定是list格式，多轮对话时，包裹[]
            padding="max_length",
            max_length=128,
            return_tensors="pt"

        )
        print(type(batch_inputs))
        print(batch_inputs.keys())
        return {
            "input_ids": batch_inputs["input_ids"],
            "attention_mask": batch_inputs["attention_mask"],
            "labels": batch_inputs["input_ids"].clone()    # 因果语言模型的labels与input_ids相同

        }   
    
    merge_temp = []
    for ds in dataset:        #
        tokenized_dataset = ds.map(preprocess, batched=True)     # 对数据集，进行批量处理
        merge_temp.append(tokenized_dataset)

    # 合并多个tokenized Dataset
    merged_tokenized_dataset = concatenate_datasets(merge_temp)


    training_args = TrainingArguments(
        output_dir="./lora",
        per_device_train_batch_size=4,      # 根据显存调整
        gradient_accumulation_steps=8,       # 梯度累积解决显存不足
        num_train_epochs=50,                  # 训练轮数
        learning_rate=3e-4,                 # 学习率
        fp16=True,                           # 混合精度训练
        logging_steps=10,
        save_strategy="steps",
        save_steps=1000,
        optim="adamw_torch")          # 优化器


    trainer = Trainer(
    model=peft_model,
    args=training_args,
    train_dataset=tokenized_dataset,
    tokenizer=tokenizer
)
    trainer.train()
    peft_model.save_pretrained("./lora_adapter")
    tokenizer.save_pretrained("./lora_adapter")


def check():
    """
    
    """
    # 加载基础模型和分词器
    base_model = "Qwen/Qwen3-1.7B"  # 替换为实际基础模型名称或路径
    tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(base_model, device_map="auto", trust_remote_code=True)

    # 加载LoRA适配器
    model = PeftModel.from_pretrained(model, "./lora_adapter")  # 替换为适配器路径
    model = model.merge_and_unload()  # 合并适配器到基础模型（可选）

    # 构建文本生成pipeline
    generator = pipeline(
        "text-generation",
        model=model,
        tokenizer=tokenizer,
        max_new_tokens=512  # 控制生成长度
    )

    outputs = generator(
        "电视剧风筝中，郑耀先为什么不策反他的兄弟？", 
        max_length=50, 
        num_return_sequences=2,
        temperature=0.7  # 控制生成随机性
    )
    print(outputs[0]["generated_text"])


    # 预训练模型的答案
    # generator = pipeline("text-generation", model="Qwen/Qwen3-1.7B", trust_remote_code=True,torch_dtype=torch.float16)  

    # outputs = generator(
    #     "张碧晨翻唱韩文歌《再见》", 
    #     max_length=100, 
    #     num_return_sequences=2,    #生成多少个回答
    #     temperature=0.5  # 控制生成随机性
    # )
    # print(outputs)



if __name__ == "__main__":
    peft_model()
    # check()

